
package com.moon.back.util;

import org.apache.commons.math3.stat.regression.OLSMultipleLinearRegression;
import java.util.ArrayList;
import java.util.List;

/**
 * ARIMA模型简单实现
 */
public class ARIMA {
    private int p;
    // 自回归项数
    private int d;
    // 差分次数
    private int q;
    // 移动平均项数
    private List<Double> data;
    private List<Double> diffData;

    public ARIMA(int p, int d, int q, List<Double> data) {
        this.p = p;
        this.d = d;
        this.q = q;
        this.data = new ArrayList<>(data);
        this.diffData = difference(data, d);
    }

    // 差分处理
    private List<Double> difference(List<Double> series, int order) {
        List<Double> diff = new ArrayList<>(series);
        for (int i = 0; i < order; i++) {
            List<Double> temp = new ArrayList<>();
            for (int j = 1; j < diff.size(); j++) {
                temp.add(diff.get(j) - diff.get(j - 1));
            }
            diff = temp;
        }
        return diff;
    }

    // 训练AR模型（简化版）
    public double[] trainAR() {
        int n = diffData.size() - p;
        double[] y = new double[n];
        double[][] x = new double[n][p];

        for (int i = p; i < diffData.size(); i++) {
            y[i - p] = diffData.get(i);
            for (int j = 0; j < p; j++) {
                x[i - p][j] = diffData.get(i - 1 - j);
            }
        }

        OLSMultipleLinearRegression regression = new OLSMultipleLinearRegression();
        regression.newSampleData(y, x);
        return regression.estimateRegressionParameters();
    }

    // 预测（简化版）
    public double forecast(int steps) {
        double[] coefficients = trainAR();
        List<Double> forecast = new ArrayList<>(diffData);

        for (int i = 0; i < steps; i++) {
            double prediction = 0;
            for (int j = 0; j < p; j++) {
                if (j < forecast.size()) {
                    prediction += coefficients[j + 1] * forecast.get(forecast.size() - 1 - j);
                }
            }
            prediction += coefficients[0];
            // 截距项
            forecast.add(prediction);
        }

        // 逆差分处理
        for (int i = 0; i < d; i++) {
            for (int j = forecast.size() - 1; j > 0; j--) {
                forecast.set(j, forecast.get(j) + forecast.get(j - 1));
            }
        }

        return forecast.get(forecast.size() - 1);
    }

    public static void main(String[] args) {
        // 示例数据
        List<Double> sampleData = List.of(
            112., 118., 132., 129., 121., 135., 
            148., 148., 136., 119., 104., 118.);

        // 创建ARIMA(1,1,0)模型
        ARIMA model = new ARIMA(1, 1, 0, sampleData);
        
        // 进行一步预测
        double prediction = model.forecast(1);
        System.out.println("Next value prediction: " + prediction);
    }
}